# #####################################################################################################################
#  批测参数配置
# -------------------------------------------
# name: 测试名称。如果为 'build' 则表示是构造数据，会尝试建表与导入数据，否则都是运行测试。
# ddl: 数据库对象类型
# primary： 主机的IP、端口
# standby： 备机的IP、端口。提供时则进行主监控备训练的业务分离，不提供则则都在主机。无备机时需要以'-'占位
# database：数据库
# username：用户名
# password：密码。注意，仅是测试工具，未提供脱敏操作，密码会在日志、屏幕等地方存在明文展示。
# pkgTypeNum：产品种类的数量
# produce-dop：生产者的并行度
# consume-dop：消费者的并行度
# runMins：运行时间
# 
# ####################################################################################################################
BATCH_CONFIG = """
# name         ddl       primary            standby            database    user    password      pkgTypeNum    produce-dop    consume-dop    runMins
# Example:
  build        common    127.0.0.1:20088    -                  postgres    carrot  dmMFkjAD      200000        20             32             30
  og-p-0050    common    127.0.0.1:20088    127.0.0.1:20098    postgres    carrot  dmMFkjAD      200000        50             32             30
"""

# 每组测试的休息间隔，秒
BATCH_TEST_INTERVAL = 5


# #####################################################################################################################
# 性能信息监控配置。 运行测试的时候，在指定的机器上运行
# #####################################################################################################################
# 开关
PERF_TRACK = False

# 机器信息, example: [(ip, user, pwd), (ip, user, pwd), ...] 
# 如果密码中有$等会被sh劫持的符号，需要用\转义，否则会登陆报错
PERF_TRACK_MACHINES = [
    ('0.0.0.0', 'carrot', 'dmMFkjAD'),
]






















# #####################################################################################################################
# 代码
# #####################################################################################################################
import os
import sys
import time
import signal
import subprocess
from datetime import datetime


class Proj:
    root_path = None
    workspace = None
    name = None
    log_file = None
    tool_path = None
    perf_track_running = False
    
    @staticmethod
    def init():
        # 1、准备路径
        Proj.root_path = os.path.abspath(os.path.dirname(sys.argv[0]))

        Proj.run_path = os.path.join(Proj.root_path, 'run')
        if not os.path.exists(Proj.run_path):
            os.system(f"mkdir -p {Proj.run_path}")
            
        Proj.name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

        Proj.workspace = os.path.join(Proj.run_path, Proj.name)
        if not os.path.exists(Proj.workspace):
            os.system(f"mkdir -p {Proj.workspace}")
            
        Proj.log_file = os.path.join(Proj.workspace, 'run.log')
        Proj.log(f'测试工作路径: {Proj.workspace}')

        # 2、编译程序
        code_path = os.path.join(Proj.root_path, 'src')
        Proj.tool_path = os.path.join(Proj.workspace, 'src')
        Proj.log(f'复制测试程序代码到: {Proj.tool_path}')

        Proj.log('开始编译测试程序')
        status = os.system(f"cp {code_path} {Proj.workspace} -r && cd {Proj.tool_path} && javac *.java")
        
        if status != 0:
            Proj.log('ERROR: 编译出错')
            Proj.end(False)
        else:
            Proj.log("编译完成")

        # 3、启动perf track
        if PERF_TRACK:
            s, o = subprocess.getstatusoutput("which sshpass")
            if s == 0:
                Proj.perf_track_running = True
            else:
                Proj.log(o)
                Proj.log("没有 sshpass，无法使用perf track")
                Proj.end(False)

        if Proj.perf_track_running:
            Proj.log("启动 perf track")
            for track_machine in PERF_TRACK_MACHINES:
                os.system(f'python3 perf/run.py install -i {track_machine[0]} -u {track_machine[1]} -w {track_machine[2]}')
                os.system(f'python3 perf/run.py run-remote -i {track_machine[0]} -u {track_machine[1]} -w {track_machine[2]} -n {Proj.name}')

    @staticmethod
    def end(normal=True):
        if Proj.perf_track_running:
            for track_machine in PERF_TRACK_MACHINES:
                os.system(f'python3 perf/run.py stop-remote -i {track_machine[0]} -u {track_machine[1]} -w {track_machine[2]} -n {Proj.name}')
                os.system(f'python3 perf/run.py download-result -i {track_machine[0]} -u {track_machine[1]} -w {track_machine[2]} -n {Proj.name}')

            os.system(f"mv perf/workspace/{Proj.name}/* {Proj.workspace}")

        Proj.log(f"所有测试数据存放在: {Proj.workspace}")
        if normal:
            Proj.log("测试结束.")
            exit(0)
        else:
            Proj.log("测试中断.")
            exit(1)


    @staticmethod
    def log(msg):
        ftime = datetime.now().strftime("%Y-%m-%d %H-%M-%S")
        with open(Proj.log_file, 'a') as f:
            f.write(msg + '\n')
        print(ftime, msg)
    
    @staticmethod
    def wait(seconds):
        if seconds <= 0:
            return
        for i in range(0, seconds):
            print(f"  等待 {i}/{seconds}s", end='\r')
            time.sleep(1)
        Proj.log(f"等待 {seconds}s, 完成                                      ")
    

class Config:
    def __init__(self, line):
        self._valid = False
        self._content = ""
        self._config_file = ""
        
        line = line.strip()
        if line == '' or line.startswith("#"):
            return
        
        params = line.split()
        if len(params) != 11:
            return
        
        self.name = params[0]
        self._config_file = os.path.join(Proj.workspace, f'{self.name}.props')
        
        self.database = params[4]
        

        self.ddl = params[1]
        self._content += f"ddl={self.ddl}\n"

        self.primary = params[2]
        self._content += f"primaryDbURL=jdbc:postgresql://{self.primary}/{self.database}?batchMode=on&loggerLevel=OFF&rewriteBatchedStatements=true\n"
        
        self.standby = params[3]
        if params[3] != '-':
            self._content += f"standbyDbURL=jdbc:postgresql://{self.standby}/{self.database}?batchMode=on&loggerLevel=OFF&rewriteBatchedStatements=true\n"
            
        self.username = params[5]
        self._content += f"username={self.username}\n"
        
        self.password = params[6]
        self._content += f"password={self.password}\n"
        
        self.machine_num = params[7]
        self._content += f"pkgTypeNum={self.machine_num}\n"
        
        self.produce_dop = params[8]
        self._content += f"produceParallel={self.produce_dop}\n"
        
        self.consume_dop = params[9]
        self._content += f"consumeParallel={self.consume_dop}\n"
        
        self.run_mins = params[10]
        self._content += f"runMins={self.run_mins}\n"

        log_file = os.path.join(Proj.workspace, f"{self.name}.res.log")
        self._content += f"logFile={log_file}\n"

        self._content += f"tablePartCount=100\n"
        
        with open(self._config_file, 'w') as f:
            f.write(self._content)
        
        self._valid = True
        
        
    def is_valid(self):
        return self._valid
    
    def __str__(self):
        return f"{self._config_file} \n{self._content}"

    def run(self):
        mode = 'build' if self.name.lower() == "build" else 'run' 
        return os.system(f"java -classpath {Proj.tool_path}:{Proj.tool_path}/postgresql.jar Main {mode} {self._config_file}")


def handle_sigint(signal, frame):
    Proj.end(False)


if __name__ == "__main__":
    signal.signal(signal.SIGINT, handle_sigint)

    Proj.init()
    Proj.log("开始批量测试")
    
    # parse config
    configs = []
    batch_config_lines = BATCH_CONFIG.split('\n')
    for line in batch_config_lines:
        config = Config(line)
        if not config.is_valid():
            continue
        configs.append(config)
    
    Proj.log(f"一共 {len(configs)} 个配置。")
    
    # run
    for i, config in enumerate(configs):
        i = i + 1
        Proj.log(f"[第 {i} 个] {config.name} #############################################################################################")
        Proj.log(f'配置文件：' + str(config))
        Proj.log('开始测试')
        config.run()
        
        Proj.log("测试结束。")

        if i < len(configs):
            Proj.wait(BATCH_TEST_INTERVAL)

    Proj.end(True)


